import cv2
import numpy as np
from sklearn.cluster import KMeans

# 读取图像
image_path = 'pic_rgb/input_data/src2.jpg'
image = cv2.imread(image_path)
if image is None:
    print(f"Error: Unable to read image at {image_path}")
    exit()

# 应用高斯模糊进行预处理，减少噪声并平滑图像
blurred_image = cv2.GaussianBlur(image, (5,5), 0)  # 使用5x5的高斯核

# 将图像从 BGR 转换为 RGB
image_rgb = cv2.cvtColor(blurred_image, cv2.COLOR_BGR2RGB)

# 将图像转换为二维数组 (height * width, 3)
image_2d = image_rgb.reshape((-1, 3))

# 使用 K-means 聚类
kmeans = KMeans(n_clusters=10, random_state=0)  # 选择聚类的数量
kmeans.fit(image_2d)

# 获取聚类中心和标签
cluster_centers = kmeans.cluster_centers_
labels = kmeans.labels_

# 将每个像素替换为对应的聚类中心颜色
segmented_image = cluster_centers[labels].reshape(image_rgb.shape).astype(np.uint8)

cv2.imshow('Original Image', image)
cv2.imshow('segmented Image', segmented_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

